from zhipuai import ZhipuAI
import json

def init_glm_ai(api_key):
    #api_key = input("智谱AI平台的API_KEY：")
    return  ZhipuAI(api_key=api_key) # 请填写您自己的APIKey

def chatTools(client, messages, tools, model_name="glm-4", temperature=0.6):
    response = client.chat.completions.create(
        model=model_name,
        messages=messages,
        tools=tools,
        tool_choice="auto",
        temperature=temperature
    )
    return response.choices[0]

def chat_glmOnce(client, messages, model_name="glm-4"):
    response = client.chat.completions.create(
        model=model_name, # 使用传入的模型名称
        messages=messages
    )
    return response.choices[0].message.content

def ask_glm(client, tools, temperature, messages, function_list, dataset, model_name="glm-4", chart_type="饼图"):
    # 注意：chart_type 参数在这里接收，但实际使用是在 ai_chat.py 的工具函数中
    ack = chatTools(client, messages, tools, model_name=model_name, temperature=temperature)

    if ack.finish_reason == "tool_calls":
        tool_call = ack.message.tool_calls[0]
        function_name = tool_call.function.name
        function_args_dict = json.loads(tool_call.function.arguments)

        # 动态注入参数
        # 确保 function_list 是一个字典，键是函数名，值是函数对象
        # 或者调整此处的逻辑以匹配您之前对 function_list 的使用方式
        # 以下假设 function_list 是一个包含工具定义的列表，如之前代码所示
        target_function = None
        for tool_def in function_list: # function_list from load_tools()
            if tool_def["function"]["name"] == function_name:
                # 获取 ai_chat.py 中定义的实际函数对象
                
                try:
                    target_function = globals()[function_name] # 这在glm.py中可能找不到ai_chat.py的函数
                                                          # 需要调整调用方式，或者将函数对象直接传入
                except KeyError:
                     return f"错误：在glm.py的全局作用域中未找到函数 {function_name}"
                break
        else:
            return "未找到匹配的工具定义"
        
   

        final_args = {
            k: v for k, v in function_args_dict.items()
        }
        final_args["dataset"] = dataset
        final_args["chart_type"] = chart_type # 将图表类型传递给工具函数

       
        if function_name == "get_city_positions":
            # 需要从 ai_chat.py 导入 get_city_positions 或将其作为参数传递
            from modules.ai_chat import get_city_positions # 临时的，不推荐的导入方式
            result = get_city_positions(**final_args)
        elif function_name == "get_average_salary":
            from modules.ai_chat import get_average_salary
            result = get_average_salary(**final_args) # avg_salary 通常不需要图表类型
        elif function_name == "get_skill_distribution":
            from modules.ai_chat import get_skill_distribution
            result = get_skill_distribution(**final_args)
        elif function_name == "get_city_specific_position":
            from modules.ai_chat import get_city_specific_position
            result = get_city_specific_position(**final_args)
        elif function_name == "get_city_salary":
            from modules.ai_chat import get_city_salary # salary 通常不需要图表类型
            result = get_city_salary(**final_args)
        else:
            return "未知的工具函数"
        return result

    return ack.message
